#include "Graphics/PiccoloVKBuffer.h"

#include <Graphics/PiccoloVKDevice.h>

namespace Piccolo
{
    PiccoloVKBuffer::PiccoloVKBuffer(PiccoloVKDevice* device, VkBufferUsageFlags usage, size_t size, void* data, bool hostViaible) : mDevice(device), mSize(size), bHostViaible(hostViaible)
    {
        CreateBuffer(usage, data);
    }

    PiccoloVKBuffer::~PiccoloVKBuffer()
    {
        VK_D(Buffer, mDevice->GetLogicDevice(), mBuffer);
        VK_F(mDevice->GetLogicDevice(), mDeviceMemory);
    }

    void PiccoloVKBuffer::CopyBuffer(PiccoloVKDevice* device, VkBuffer srcBuffer, VkBuffer dstBuffer, size_t size)
    {
        VkCommandBuffer commandBuffer = device->CreateAndBeginSingleCommandBuffer();

        VkBufferCopy bufferCopy =
        {
            .srcOffset = 0,
            .dstOffset = 0,
            .size = size
        };
        vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, &bufferCopy);

        device->EndSingleCommandBuffer(commandBuffer);
    }

    void PiccoloVKBuffer::CreateBufferInternal(PiccoloVKDevice* device, VkMemoryPropertyFlags memProps, VkBufferUsageFlags usage, size_t size, VkBuffer* outBuffer, VkDeviceMemory* outMemory)
    {
        VkBufferCreateInfo bufferInfo =
        {
            .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .size = size,
            .usage = usage,
            .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
            .queueFamilyIndexCount = 0,
            .pQueueFamilyIndices = nullptr
        };
        CALL_VK(vkCreateBuffer(device->GetLogicDevice(), &bufferInfo, nullptr, outBuffer));

        // 查询内存信息
        VkMemoryRequirements memRequirements;
        vkGetBufferMemoryRequirements(device->GetLogicDevice(), *outBuffer, &memRequirements);
        // 分配VkBuffer内存
        VkMemoryAllocateInfo allocateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
            .pNext = nullptr,
            .allocationSize = memRequirements.size,
            .memoryTypeIndex = static_cast<uint32_t>(device->GetMemoryIndex(memProps, memRequirements.memoryTypeBits))
        };
        CALL_VK(vkAllocateMemory(device->GetLogicDevice(), &allocateInfo, nullptr, outMemory));
        CALL_VK(vkBindBufferMemory(device->GetLogicDevice(), *outBuffer, *outMemory, 0));
    }

    /**
     * 如果CPU可以访问该buffer，创建buffer后直接map映射
     * 如果CPU无法访问，需要创建一个中转buffer再拷贝数据
     * @param usage 用途
     * @param data 数据
     */
    void PiccoloVKBuffer::CreateBuffer(VkBufferUsageFlags usage, void* data)
    {
        if (bHostViaible)
        {
            // CPU可以访问
            CreateBufferInternal(mDevice, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, usage, mSize, &mBuffer, &mDeviceMemory);
            WriteData(data);
        }
        else
        {
            // CPU不能访问
            VkBuffer stageBuffer;
            VkDeviceMemory stageMemory;
            CreateBufferInternal(mDevice, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, mSize, &stageBuffer, &stageMemory);

            void* mapping;
            CALL_VK(vkMapMemory(mDevice->GetLogicDevice(), stageMemory, 0, VK_WHOLE_SIZE, 0, &mapping));
            memcpy(mapping, data, mSize);
            vkUnmapMemory(mDevice->GetLogicDevice(), stageMemory);

            CreateBufferInternal(mDevice, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, usage | VK_BUFFER_USAGE_TRANSFER_DST_BIT, mSize, &mBuffer, &mDeviceMemory);

            CopyBuffer(mDevice, stageBuffer, mBuffer, mSize);

            // 释放临时buffer的内存
            VK_D(Buffer, mDevice->GetLogicDevice(), stageBuffer);
            VK_F(mDevice->GetLogicDevice(), stageMemory);
        }
    }

    /**
     * CPU可以访问的，就可以往buffer写数据
     * @param data 数据
     * @return VkResult
     */
    VkResult PiccoloVKBuffer::WriteData(void *data)
    {
        if (data && bHostViaible)
        {
            void* mapping;
            VkResult ret = vkMapMemory(mDevice->GetLogicDevice(), mDeviceMemory, 0, VK_WHOLE_SIZE, 0, &mapping);
            memcpy(mapping, data, mSize);
            vkUnmapMemory(mDevice->GetLogicDevice(), mDeviceMemory);
            return ret;
        }
        return VK_ERROR_INITIALIZATION_FAILED;
    }
}
